import os
from argparse import ArgumentParser
import numpy as np
import torch
from PIL import Image

from Models_GAN import Generator, Generator_DC, Discriminator
import Dataloader_Uv

CKPT_DIR = 'models'
G_FN = 'gan_g.pth'
MAX_SEQ_LEN = 100
FILENAME = 'sample.jpeg'
SPLIT = '/'

def save_data(filename,pic_data):
    picdata = pic_data.reshape([100,64,64])
    v = np.hstack(picdata).reshape(64, 6400)
    v = v * 255
    im = Image.fromarray(v)
    if im.mode == "F":
        im = im.convert('L')        
    im.save("Generator_Out" + SPLIT + filename + ".jpeg")

def generate(arg):
    ''' Sample MIDI from trained generator model
    '''
    # prepare model
    if arg.generator_DC:
        num_feats = 100
    else:
        num_feats = 64*64

    use_gpu = torch.cuda.is_available()
    g_model = Generator(num_feats, use_cuda=use_gpu)

    if not use_gpu:
        ckpt = torch.load(os.path.join(CKPT_DIR, G_FN), map_location='cpu')
    else:
        ckpt = torch.load(os.path.join(CKPT_DIR, G_FN))

    g_model.load_state_dict(ckpt)

    # generate from model then save 
    g_states = g_model.init_hidden(1)   #init_hidden()的参数为real_batch_size 这里默认为1
    z = torch.empty([1, MAX_SEQ_LEN, num_feats]).uniform_() # random vector
    if use_gpu:
        z = z.cuda()
        g_model.cuda()

    g_model.eval()

    pic_list_data = []
    for i in range(arg.n):#将多个序列进行拼接    这里还是默认一次生成一份100张图片的序列
        g_feats, g_states = g_model(z, g_states)
        pic_data = g_feats.squeeze().cpu()
        pic_data = pic_data.detach().numpy() 
        pic_list_data.append(pic_data)

    if len(pic_list_data) > 1:
        pic_list_data = np.concatenate(pic_list_data, axis=0)
    else:
        pic_list_data = pic_list_data[0]

    save_data(FILENAME, pic_list_data)
    # print('Generated {}'.format(FILENAME))



if __name__ == "__main__":
    ARG_PARSER = ArgumentParser()
    # number of times to execute generator model;
    # all generated data are concatenated to form a single longer sequence
    ARG_PARSER.add_argument('--n', default=1, type=int)
    ARG_PARSER.add_argument('--filename', default='sample', type=str)
    ARG_PARSER.add_argument('--generator_DC', action='store_true')                          #替换生成网络结构   默认为否
    if(platform.system()=='Windows'):
        SPLIT = '\\'
    else:
        SPLIT = '/'
    ARGS = ARG_PARSER.parse_args()
    if ARGS.generator_DC:
        G_FN = 'gan_DC_g.pth'
    FILENAME = ARGS.filename
    print(FILENAME)
    generate(ARGS)
